package com.xjh.serverinspect.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.channel.ClientChannel;
import org.apache.sshd.client.channel.ClientChannelEvent;
import org.apache.sshd.client.session.ClientSession;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * <p>
 * sshd执行器
 * </p>
 *
 * @author xujinghui
 * @since 2024-10-14
 */
@Slf4j
public class SshdExecutor {

    private SshClient client;
    private ClientSession session;
    private boolean connected;

    /**
     * 密码方式
     *
     * @param ip       服务IP
     * @param port     服务端口
     * @param user     登录用户名
     * @param password 登录密码
     */
    public SshdExecutor(String ip, Integer port, String user, String password) {
        client = SshClient.setUpDefaultClient();
        client.start();
        try {
            session = client.connect(user, ip, port).verify(10 * 1000).getSession();
            if (password != null && !password.isEmpty()) {
                session.addPasswordIdentity(password);
            }
            if (session.auth().verify(10 * 1000).isSuccess()) {
                connected = true;
            } else {
                log.error("auth failed");
            }
        } catch (IOException e) {
            log.error("创建ssh连接失败", e);
        }
    }

    /**
     * 执行命令
     *
     * @param command 命令字符串
     */
    public String[] executeCommand(String command) {
        if (connected) {
            try(ClientChannel channel = session.createExecChannel(command)) {
                ByteArrayOutputStream out = new ByteArrayOutputStream();
                ByteArrayOutputStream err = new ByteArrayOutputStream();
                channel.setOut(out);
                channel.setErr(err);

                if (!channel.open().verify(10 * 1000).isOpened()) {
                    throw new Exception("open failed");
                }
                List<ClientChannelEvent> list = new ArrayList<>();
                list.add(ClientChannelEvent.CLOSED);
                channel.waitFor(list, 10 * 1000);
                return new String[]{out.toString(), err.toString()};
            } catch (Exception e) {
                log.error("执行命令失败", e);
                return new String[]{null, e.getMessage()};
            }
        } else {
            return new String[]{null, "未连接到服务器"};
        }
    }

    /**
     * 执行命令并关闭连接
     *
     * @param command 命令字符串
     * @return
     */
    public String[] executeCommandAndClose(String command) {
        String[] result = executeCommand(command);
        close();
        return result;
    }

    public void close() {
        if (client != null) {
            try {
                client.stop();
                client.close();
            } catch (IOException e) {
                log.error("关闭ssh连接失败", e);
            }
        }
    }

}
